from octis.evaluation_metrics.diversity_metrics import TopicDiversity
from octis.dataset.dataset import Dataset
from octis.evaluation_metrics.coherence_metrics import Coherence
from octis.models.GINOPIC import GINOPIC

if __name__ == "__main__":
  dataset = Dataset()
  dataset_name = "20NewsGroup"
  # dataset.fetch_dataset("20NewsGroup")
  dataset.load_custom_dataset_from_folder("../preprocessed_datasets/" + dataset_name)
  k = 10

  # True 则模型将在训练集和测试集上的评估(默认值：True)
  partition = True
  # 验证
  validation = False

  # 20NGS
  params = {
    'num_gin_layers': 2,
    'g_feat_size': 2048,
    'num_mlp_layers': 1,
    'gin_hidden_dim': 200,
    'gin_output_dim': 768,
    'eps_simGraph': 0.4
  }

  # BBC
  # params ={
  #   'num_gin_layers': 3,
  #   'g_feat_size': 256,
  #   'num_mlp_layers': 1,
  #   'gin_hidden_dim': 50,
  #   'gin_output_dim': 512,
  #   'eps_simGraph': 0.3
  # }

  model = GINOPIC(num_topics=k,
                  use_partitions=partition, # TRUE 数据集划分
                  use_validation=validation, # FAlSE
                  num_epochs=50,
                  w2v_path='./w2v/{}_part{}_valid{}/'.format(dataset_name, partition, validation),
                  graph_path='./doc_graphs/{}_part{}_valid{}/'.format(dataset_name, partition, validation),
                  num_gin_layers=params['num_gin_layers'],
                  g_feat_size=params['g_feat_size'],
                  num_mlp_layers=params['num_mlp_layers'],
                  gin_hidden_dim=params['gin_hidden_dim'],
                  gin_output_dim=params['gin_output_dim'],
                  eps_simGraph=params['eps_simGraph'],
                  seed=10
                  )
  # model = CTM(num_topics=10, num_epochs=30, inference_type='zeroshot', bert_model="bert-base-nli-mean-tokens")
  # model = ETM(num_topics=20, num_epochs=100)
  # model = GNTM(num_topics=20, num_epochs=100)

  output = model.train_model(dataset)
  for t in output['topics'][:5]:
    print(" ".join(t))

  # Initialize metric
  CV = Coherence(texts=dataset.get_corpus(), topk=10, measure='c_v')
  npmi = Coherence(texts=dataset.get_corpus(), topk=10, measure='c_npmi')
  print("NPMI: " + str(npmi.score(output)))
  print("Cv: " + str(CV.score(output)))

  # Retrieve metrics score
  # topic_diversity = TopicDiversity(topk=10)
  # topic_diversity_score = topic_diversity.score(output)
  # print("Topic diversity: " + str(topic_diversity_score))

